# coding: utf-8
import os
import random
import jieba
import nltk
import sklearn
from sklearn.naive_bayes import MultinomialNB
import numpy as np
import matplotlib.pyplot as plt


# 词去重
def filter_word_set(words_file):
    words_set = set()
    with open(words_file, 'r', encoding='UTF-8') as fp:
        for line in fp.readlines():
            word = line.strip()
            if len(word) > 0 and word not in words_set:
                words_set.add(word)
    return words_set


# 文本处理，样本生成过程
def text_process(folder_path_var, test_size=0.2):
    folder_list = os.listdir(folder_path_var)
    data_list = []
    class_list = []

    # 遍历文件夹 # 读取文件
    for folder in folder_list:
        new_folder_path = os.path.join(folder_path_var, folder)
        files = os.listdir(new_folder_path)

        j = 1
        for file in files:
            if j > 100:
                break
            with open(os.path.join(new_folder_path, file), 'r', encoding='UTF-8', errors='ignore') as fp:
                raw = fp.read()
            word_cut = jieba.cut(raw, cut_all=False)
            word_list = list(word_cut)

            data_list.append(word_list)
            class_list.append(folder.encode('utf-8'))
            j += 1

    # 划分训练集和测试集
    data_class_list = list(zip(data_list, class_list))
    random.shuffle(data_class_list)
    index = int(len(data_class_list) * test_size) + 1
    train_list = data_class_list[index:]
    test_list = data_class_list[:index]
    train_data_list_var, train_class_list_var = zip(*train_list)
    test_data_list_var, test_class_list_var = zip(*test_list)

    # 统计词频记入all_words_dict
    all_words_dict = {}
    for word_list in train_data_list_var:
        for word in word_list:
            if all_words_dict.get(word) is not None:
                all_words_dict[word] += 1
            else:
                all_words_dict[word] = 1

    # 按照词频进行降序排序
    all_words_tuple_list = sorted(all_words_dict.items(), key=lambda f: f[1], reverse=True)
    all_words_list_var = list(list(zip(*all_words_tuple_list))[0])

    return all_words_list_var, train_data_list_var, test_data_list_var, train_class_list_var, test_class_list_var


def words_dict(all_words_list_var, deleteN, stopwords_set_var=None):
    # 选取特征词
    if stopwords_set_var is None:
        stopwords_set_var = set()
    feature_words_var = []
    n = 1
    for t in range(deleteN, len(all_words_list_var), 1):
        if n > 1000:   # 特征词的维度1000
            break

        if not all_words_list_var[t].isdigit() and all_words_list_var[t] not in stopwords_set_var and 1 < len(
                all_words_list_var[t]) < 5:
            feature_words_var.append(all_words_list_var[t])
            n += 1
    return feature_words_var


# 文本特征
def text_features(train_data_list, test_data_list, feature_words, flag='nltk'):
    def text_features(text, feature_words):
        text_words = set(text)

        if flag == 'nltk':
            ## nltk特征 dict
            features = {word: 1 if word in text_words else 0 for word in feature_words}
        elif flag == 'sklearn':
            ## sklearn特征 list
            features = [1 if word in text_words else 0 for word in feature_words]
        else:
            features = []

        return features

    train_feature_list = [text_features(text, feature_words) for text in train_data_list]
    test_feature_list = [text_features(text, feature_words) for text in test_data_list]
    return train_feature_list, test_feature_list


# 分类，同时输出准确率等
def text_classifier(train_feature_list, test_feature_list, train_class_list, test_class_list, flag='nltk'):

    if flag == 'nltk':
        # 使用nltk分类器
        # print("wdnmd")
        train_flist = zip(train_feature_list, train_class_list)
        test_flist = zip(test_feature_list, test_class_list)
        classifier = nltk.classify.NaiveBayesClassifier.train(train_flist)
        test_accuracy_var = nltk.classify.accuracy(classifier, test_flist)

    elif flag == 'sklearn':
        # sklearn分类器
        classifier = MultinomialNB().fit(train_feature_list, train_class_list)
        test_accuracy_var = classifier.score(test_feature_list, test_class_list)
    else:
        test_accuracy_var = []
    print(test_accuracy_var)
    return test_accuracy_var


print("start")

# 文本预处理
folder_path = './20_newsgroups'
all_words_list, train_data_list, test_data_list, train_class_list, test_class_list = text_process(folder_path,
                                                                                                  test_size=0.2)

# 生成stopwords_set
stopwords_file = './stopwords_cn.txt'
stopwords_set = filter_word_set(stopwords_file)

# 文本特征提取和分类
# flag = 'nltk'
flag = 'sklearn'
size = range(0, 1000, 20)
test_accuracy_list = []
for size1 in size:
    feature_words = words_dict(all_words_list, size1, stopwords_set)
    train_feature_list, test_feature_list = text_features(train_data_list, test_data_list, feature_words, flag)
    test_accuracy = text_classifier(train_feature_list, test_feature_list, train_class_list, test_class_list, flag)
    test_accuracy_list.append(test_accuracy)
print(test_accuracy_list)



# 结果评价
plt.figure()
plt.plot(size, test_accuracy_list)
plt.title('Relationship of deleteNs and test_accuracy')
plt.xlabel('deleteNs')
plt.ylabel('test_accuracy')
plt.show()
plt.savefig('result.png')

print("finished")
